package com.heima.datastructure.avltree;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * <strong>AVL树</strong>
 * <ul>
 *     <li>二叉搜索树在插入和删除时，节点可能失衡</li>
 *     <ul>
 *         <li>如果一个节点的左右孩子，高度差超过1，则此节点失衡</li>
 *         <li>需要旋转恢复平衡</li>
 *     </ul>
 *     <li>如果在插入和删除时通过旋转，始终让二叉搜索树保持平衡，称为自平衡的二叉搜索树</li>
 *     <li>AVL是自平衡二叉搜索树的实现之一</li>
 * </ul>
 *
 * @author 勾新杰
 * @version 1.0
 * @date 2024/9/13 16:29
 */
public class AVLTree<K extends Comparable<K>, V> {

    private AVLNode<K, V> root; // AVL二叉搜索树的根节点，root默认初始化null

    /**
     * 无参构造，root初始默认就是null
     */
    public AVLTree() {
    }

    /**
     * 有参构造，root为传入的未做平衡处理的二叉搜索树的根节点
     *
     * @param root 根节点
     */
    AVLTree(AVLNode<K, V> root) {
        this.root = balanceTree(root);
    }

    /**
     * AVL节点类
     *
     * @param <K> key的数据类型，要实现Comparable接口
     * @param <V> value的数据类型，存放数据
     */
    static class AVLNode<K, V> {
        K key; // key
        V value; // value
        AVLNode<K, V> left; // 左子节点
        AVLNode<K, V> right; // 右子节点
        Integer height = 1; // 当前节点的高度，默认设置为1，这里的height是左子树的高度和右子树的高度的最大值加一，叶子节点高度都是1

        public AVLNode(K key) {
            this.key = key;
        }

        public AVLNode(K key, V value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(K key, V value, Integer height) { // 测试用，实际用不到
            this.key = key;
            this.value = value;
            this.height = height;
        }

        public AVLNode(K key, V value, AVLNode<K, V> left, AVLNode<K, V> right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }

        public AVLNode(K key, V value, AVLNode<K, V> left, AVLNode<K, V> right, Integer height) { // 测试用
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            this.height = height;
        }
    }

    /**
     * 求一个节点的高度
     *
     * @param node 节点
     * @return 该节点的高度
     */
    private Integer height(AVLNode<K, V> node) {
        return node == null ? 0 : node.height;
    }

    /**
     * 更新节点高度
     *
     * @param node 节点
     */
    private void updateHeight(AVLNode<K, V> node) {
        node.height = Integer.max(height(node.left), height(node.right)) + 1;
    }

    // 平衡因子（balance factor）= 左子树高度 - 右子树高度

    /**
     * 求一个节点的平衡因子
     * <ul>
     *     <li>返回值为0，1，-1 表示平衡</li>
     *     <li>返回值 > 1 或 返回值 < -1 表示不平衡</li>
     * </ul>
     *
     * @param node 节点
     * @return 该节点的平衡因子
     */
    private Integer balanceFactor(AVLNode<K, V> node) {
        return height(node.left) - height(node.right);
    }

    // 一、LL
    //    1、失衡节点的平衡因子 > 1，也就是左边更高
    //    2、失衡节点的左孩子的平衡因子 >= 0，也就是左孩子也是左边更高或者等高
    //    3、需要进行一次右旋达到平衡，R
    // 二、LR
    //    1、失衡节点的平衡因子 > 1，也就是左边更高
    //    2、失衡节点的左孩子的平衡因子 < 0，也就是左孩子是右边更高
    //    3、需要进行一次左旋，变成LL的情况，再进行一次右旋，达到平衡，LR
    // 三、RL
    //    1、失衡节点的平衡因子 < -1，也就是右边更高
    //    2、失衡节点的右孩子的平衡因子 > 0，也就是右孩子是左边更高
    //    3、需要进行一次右旋，变成RR的情况，再进行一次左旋，达到平衡，RL
    // 四、RR
    //    1、失衡节点的平衡因子 < -1，也就是右边更高
    //    2、失衡节点的右孩子的平衡因子 <= 0，也就是右孩子也是右边更高或者等高
    //    3、需要进行一次左旋达到平衡，L

    /**
     * 右旋操作R => 解决LL
     *
     * @param root 要右旋的树的根节点
     * @return 右旋后该树的根节点
     */
    private AVLNode<K, V> rightRotate(AVLNode<K, V> root) {
        AVLNode<K, V> left = root.left;
        AVLNode<K, V> lRight = left.right; // 不用考虑left为空指针的情况，因为右旋就是处理左边高的情况
        left.right = root;
        root.left = lRight; // 第三行和第四行可以交换，但是第三行不能和第二行交换，会覆盖
        updateHeight(root); // 先更新换下来的节点
        updateHeight(left); // 再更新换上去的节点
        return left;
    }

    /**
     * 左旋操作L =>解决RR
     *
     * @param root 要左旋的树的根节点
     * @return 左旋后该树的根节点
     */
    private AVLNode<K, V> leftRotate(AVLNode<K, V> root) {
        AVLNode<K, V> right = root.right;
        AVLNode<K, V> rLeft = right.left;
        right.left = root;
        root.right = rLeft;
        updateHeight(root);
        updateHeight(right);
        return right;
    }

    /**
     * 左右旋操作LR => 解决LR，先左旋当前树的左子树，旋转好的左子树接到当前树root根节点，再右旋当前树
     *
     * @param root 要左右旋的树的根节点
     * @return 旋转后该树的根节点
     */
    private AVLNode<K, V> leftRightRotate(AVLNode<K, V> root) {
        root.left = leftRotate(root.left);
        return rightRotate(root);
    }

    /**
     * 右左旋操作RL => 解决RL：先右旋当前树的右子树，旋转好的右子树接到当前树的root根节点，再左旋当前树
     *
     * @param root 要右左旋的树的根节点
     * @return 旋转后该数的根节点
     */
    private AVLNode<K, V> rightLeftRotate(AVLNode<K, V> root) {
        root.right = rightRotate(root.right);
        return leftRotate(root);
    }

    /**
     * 平衡一颗搜索二叉树或子树的根节点，平衡一个节点
     *
     * @param root 根节点
     * @return 平衡后的树的根节点
     */
    private AVLNode<K, V> balance(AVLNode<K, V> root) {
        if (root == null) {
            return null;
        }
        int balanceFactor = balanceFactor(root);
        if (balanceFactor > 1 && balanceFactor(root.left) >= 0) { // LL
            return rightRotate(root);
        } else if (balanceFactor > 1 && balanceFactor(root.left) < 0) { // LR
            return leftRightRotate(root);
        } else if (balanceFactor < -1 && balanceFactor(root.right) > 0) { // RL
            return rightLeftRotate(root);
        } else if (balanceFactor < -1 && balanceFactor(root.right) <= 0) { // RR
            return leftRotate(root);
        } else {
            return root;
        }
    }

    /**
     * 平衡一颗二叉搜索树
     *
     * @param root 根节点
     * @return 平衡后的树的根节点
     */
    private AVLNode<K, V> balanceTree(AVLNode<K, V> root) {
        if (root == null) {
            return null;
        }
        root.left = balanceTree(root.left);
        root.right = balanceTree(root.right);
        return balance(root);
    }

    /**
     * put方法，实现往二叉搜索树中插入节点，并平衡插入后的二叉树
     *
     * @param key   key
     * @param value value
     */
    public void put(K key, V value) {
        root = doPut1(root, key, value);
        // root = doPut2(key,value);
    }

    /**
     * 基于递归方式实现put方法
     *
     * @param root  当前树的根节点
     * @param key   key
     * @param value value
     * @return 当前树添加指定节点后并平衡的二叉搜索树的根节点
     */
    private AVLNode<K, V> doPut1(AVLNode<K, V> root, K key, V value) {
        // 递归结束条件：找到空位添加，也就是root == null
        if (root == null) {
            return new AVLNode<>(key, value);
        }
        // 如果key要在当前树的左子树添加，则递归调用左子树添加
        if (key.compareTo(root.key) < 0) {
            root.left = doPut1(root.left, key, value);
        }
        // 如果key要在当前树的右子树添加，则递归调用右子树添加
        else if (root.key.compareTo(key) < 0) {
            root.right = doPut1(root.right, key, value);
        }
        // 如果key存在于树中的某个节点，则改变该值
        else {
            root.value = value;
            return root;
        }
        // 添加完了过后可能不是一个平衡二叉树，而且高度也会变化
        updateHeight(root);
        return balance(root); // 归的时候检查并平衡当前树
    }

    /**
     * 用非递归方式也就是迭代方式实现put方法
     *
     * @param key   key
     * @param value value
     * @return 插入节点重构后二叉搜索树的根节点
     */
    private AVLNode<K, V> doPut2(K key, V value) {
        AVLNode<K, V> curr = root;
        AVLNode<K, V> parent = null; // 记录父节点
        LinkedList<AVLNode<K, V>> stack = new LinkedList<>(); // 记录来时路，回去的时候这条路上的节点都要更新高度以及平衡，一级一级的往上平衡
        while (curr != null) {
            if (key.compareTo(curr.key) < 0) { // key在左边插入
                parent = curr;
                stack.push(curr);
                curr = curr.left;
            } else if (curr.key.compareTo(key) < 0) { // key在右边插入
                parent = curr;
                stack.push(curr);
                curr = curr.right;
            } else { // 已存在key
                curr.value = value;
                return root;
            }
        }
        AVLNode<K, V> node = new AVLNode<>(key, value);
        if (parent == null) { // 父节点没有，说明原本树为空
            return node;
        } else if (key.compareTo(parent.key) < 0) { // 待插入key在其父节点左边，接上去
            parent.left = node;
        } else { // 待插入key在其父节点右边，接上去
            parent.right = node;
        }
        while (!stack.isEmpty()) { // 来时的路依次弹出栈，进行更新高度和平衡操作
            AVLNode<K, V> last = stack.pop(); // 第一次弹出栈的就是parent
            updateHeight(last); // 第一次节点已经接上parent了，直接更新高度
            parent = balance(last); // 第一次里面，parent就是当前子树的根节点，平衡它
            if (!stack.isEmpty() && last.key.compareTo(stack.peek().key) < 0) { // 栈非空而且parent在其父节点的左边
                stack.peek().left = parent; // 则将当前的子树接到根节点parent的父节点上
            } else if (!stack.isEmpty() && stack.peek().key.compareTo(last.key) < 0) { // 栈非空而且parent在其父节点的右边
                stack.peek().right = parent; // 则将当前的子树接到根节点parent的父节点上
            }
        }
        return parent; // parent就是一层一层上上走平衡后的子树的根节点，走到最后就是整棵树的根节点
    }

    /**
     * remove方法，删除指定节点，并重构二叉搜索树保持平衡以及不破坏特性
     * @param key key
     * @return 删除节点的值
     */
    public V remove(K key) {
        List<V> result = new ArrayList<>();
        root = doRemove1(root, key, result);
        return result.get(0);
        // return doRemove2(key);
    }

    /**
     * 递归方式实现remove方法
     *
     * @param root   当前树的根节点
     * @param key    要删除的节点的key
     * @param result 用来存放删除节点值得list集合，因为在递归里面得传可变类型
     * @return 在当前树中删除该节点重构平衡后的根节点
     */
    private AVLNode<K, V> doRemove1(AVLNode<K, V> root, K key, List<V> result) {
        // 特殊情况，当前树为空，删不了一点，也包括树里面没有该节点
        if (root == null) {
            return null;
        }
        if (key.compareTo(root.key) < 0) { // 待删除节点在当前树的左子树
            root.left = doRemove1(root.left, key, result);
        } else if (root.key.compareTo(key) < 0) { // 待删除节点在当前树的右子树
            root.right = doRemove1(root.right, key, result);
        } else { // 待删除节点就是当前树的根节点
            result.add(root.value); // 存取删除节点值的value
            if (root.left == null && root.right == null) { // 左子树右子树都为null，要单独处理是因为，它没有子树，处理高度的时候会报空指针
                return null;
            } else if (root.left == null) { // 待删除节点的左子树为null，右子树不为null
                root = root.right; // 则将右子树作为删除后的新树，并修改高度以及平衡操作
            } else if (root.right == null) { // 待删除节点的右子树为null，左子树不为null
                root = root.left; // 则将左子树作为删除后的新树，并修改高度以及平衡操作
            } else { // 左子树和右子树都不为null
                AVLNode<K, V> successor = root.right;
                while (successor.left != null) { // 必有右子树，找到后继节点
                    successor = successor.left;
                }
                successor.right = doRemove1(root.right, successor.key, new ArrayList<>()); // 在待删除节点的右子树中删除后继节点，并把删除后的右子树接到root根节点，再让后继节点去替换root节点，即将root的左右子树接到后继节点
                successor.left = root.left;
                root = successor; // 赋值给root，好统一操作
            }
        }
        updateHeight(root); // 统一更新高度
        return balance(root); // 统一平衡搜索二叉树，并返回平衡后的根节点
    }

    /**
     * 非递归方式实现remove方法，迭代，删除平衡二叉搜索树的某个节点
     * @param key key
     * @return 删除该节点重构平衡后该树的根节点
     */
    private V doRemove2(K key) {
        AVLNode<K, V> curr = root;
        AVLNode<K, V> parent = null;
        LinkedList<AVLNode<K, V>> stack = new LinkedList<>();
        while (curr != null) {
            if (key.compareTo(curr.key) < 0) {
                parent = curr;
                stack.push(curr);
                curr = curr.left;
            } else if (curr.key.compareTo(key) < 0) {
                parent = curr;
                stack.push(curr);
                curr = curr.right;
            } else {
                break;
            }
        }
        if (curr == null) {
            return null;
        }
        if (curr.left == null && curr.right == null) {
            shift(parent, curr, null);
        } else if (curr.left == null) {
            shift(parent, curr, curr.right);
        } else if (curr.right == null) {
            shift(parent, curr, curr.left);
        } else {
            AVLNode<K, V> successor = curr.right;
            AVLNode<K, V> sParent = null;
            while (successor.left != null) {
                sParent = successor;
                stack.push(successor);
                successor = successor.left;
            }
            if (curr.right == successor) {
                shift(parent, curr, successor);
                successor.left = curr.left;
            } else {
                shift(sParent, successor, successor.right);
                shift(parent, curr, successor);
                successor.left = curr.left;
                successor.right = curr.right;
            }
            while (stack.peek() != parent){
                AVLNode<K, V> last = stack.pop();
                updateHeight(last);
                sParent = balance(last);
                if (!stack.isEmpty() && stack.peek() != parent && stack.peek().left == last){
                    stack.peek().left = sParent;
                } else if (!stack.isEmpty() && stack.peek() != parent && stack.peek().right == last) {
                    stack.peek().right = sParent;
                }
            }
            successor.right = sParent;
        }
        while (!stack.isEmpty()){
            AVLNode<K, V> last = stack.pop();
            updateHeight(last);
            parent = balance(last);
            if (!stack.isEmpty() && stack.peek().left == last){
                stack.peek().left = parent;
            } else if (!stack.isEmpty() && stack.peek().right == last) {
                stack.peek().right = parent;
            }
        }
        if (parent != null) {
            root = parent;
        }
        return curr.value;
    }

    /**
     * 托孤方法
     * @param parent 要托孤给谁，delete的父亲
     * @param delete 删除delete，让child顶替他的位置
     * @param child 需要托孤的节点，不一定是delete的子节点
     */
    private void shift(AVLNode<K, V> parent, AVLNode<K, V> delete, AVLNode<K, V> child) {
        if (parent == null) {
            root = child;
        } else if (parent.left == delete) {
            parent.left = child;
        } else if (parent.right == delete) {
            parent.right = child;
        }
    }

    /**
     * 判断当前树与另一颗AVL二叉搜索树是否是同一颗树
     *
     * @param tree 另一颗AVL二叉树
     * @return 是同一颗返回true
     */
    public boolean isSameTree(AVLTree<K, V> tree) {
        return isSameTree(this.root, tree.root); // 还处于这个类里面，可以.root方式获取
    }

    /**
     * 判断当前树与另外一颗未做平衡处理的普通二叉搜索树是否一样
     *
     * @param root 另一颗普通二叉搜索树的根节点
     * @return 一样返回true
     */
    boolean isSameTree(AVLNode<K, V> root) { // 不加public防止AVLNode爆黄，范围是当前包
        return isSameTree(this.root, root);
    }

    /**
     * 递归实现的比较两颗二叉搜索树是否一样
     *
     * @param root1 树1
     * @param root2 树2
     * @return 一样返回true
     */
    private boolean isSameTree(AVLNode<K, V> root1, AVLNode<K, V> root2) {
        if (root1 == null && root2 == null) {
            return true;
        } else if (root1 == null || root2 == null) {
            return false;
        } else {
            if (root1.key.compareTo(root2.key) == 0) {
                if (root1.value.equals(root2.value)) {
                    return isSameTree(root1.left, root2.left) && isSameTree(root1.right, root2.right);
                }
                return false;
            }
            return false;
        }
    }
}
